fix(qwen3_5_moe): detect shared-expert quant for modelopt mixed checkpoints - #294

Fixes #164
Fixes #183
Related: #274 (same shared-expert crash family on another mixed-quant uncensored checkpoint; the dense side routes through #296, but that exact checkpoint is untested)

Problem

apodex/Apodex-1.1-mini-NVFP4 (Qwen3.5 MoE A3B, 256 experts) is an nvidia modelopt MIXED_PRECISION checkpoint: the routed experts are NVFP4, but the shared expert and attn/GDN projections are per-tensor FP8 (quant_algo: "FP8" in quantized_layers, fp8-e4m3 weight + scalar weight_scale + input_scale, no weight_scale_2).

parse_config assumed the shared expert always matched the routed experts:

dense_quant = "nvfp4" if expert_quant == "nvfp4" else _dense_mlp_quant(hf_config)

So dense_quant was forced to "nvfp4", the shared expert was built as Nvfp4DenseColMerged / Nvfp4DenseLinear, and weight loading crashed with

KeyError: 'model.layers.0.mlp.shared_expert.gate_up_proj.weight'
  at freetoken/kernel/triton/nvfp4_linear.py:866

because an FP8 shared expert has no weight_scale_2 / weight_global buffers. vLLM/SGLang serve the same checkpoint fine, so this is purely the FreeToken mixed-precision path.

Changes

models/qwen3_5_moe/config.py

• Add _shared_expert_quant(): probe the modelopt quantized_layers map for the shared-expert quant_algo. Returns "nvfp4" when packed FP4 (kept native W4A16) or "none" when FP8/bf16 (dequantized to bf16 at load). Pure-NVFP4 checkpoints have no per-layer map and keep the native FP4 default; dense (num_experts == 0) checkpoints are unaffected.

models/qwen3_5_moe/weight.py

• In _iter_weights_attn_fp8, the native fp8 (W8A16) path now only applies to attn/GDN projections (.self_attn. / .linear_attn.), which actually have fp8 linears. Any other per-tensor FP8 weight — the MoE shared expert on mixed checkpoints — falls through to the existing bf16 dequant path, where the already-present shared-expert gate/up fusion (shared_expert.{gate,up}_proj -> gate_up_proj) builds the bf16 state dict.

tests/models/test_qwen3_5_moe_shared_expert_quant.py

• Unit tests for _shared_expert_quant: no quant config, NVFP4 without layer map, mixed with FP8 shared expert, mixed with NVFP4 shared expert, and a map without a shared-expert entry. The FP8 case fails on main (the detection is absent and the checkpoint crashes at load) and passes with this change.

Pure NVFP4, block-FP8, dense NVFP4, and compressed-tensors NVFP4 checkpoints take unchanged code paths.

Tested on (real hardware)

• GPU: NVIDIA GeForce RTX 3080 Laptop, 16 GB VRAM (sm_86), driver 610.57.04
• CPU: Intel Core i7-10870H (8 cores), 32 GB RAM + 31 GB swap
• OS: CachyOS (Linux, x86_64)
• FreeToken 0.1.2, torch 2.11.0+cu130 (CUDA 13.0)
• Tested commit: e6e2b26 on top of main (4b94bdc)
• Checkpoint: apodex/Apodex-1.1-mini-NVFP4
• Command: ft serve --model-path <Apodex-1.1-mini-NVFP4> --moe-backend offload --moe-cache-auto --port 8899 --host 127.0.0.1
• main: crashes at weight load with the KeyError above.
• This branch: model loads, API server is ready, /v1/chat/completions returns correct generations; decode throughput 44–60 tok/s in the engine logs (gen throughput (token/s) lines); GPU utilization shows the expected MoE-offload bursts (~87–100% during active decode, power up to ~114 W vs ~34 W idle); ~14.6/16 GB VRAM used.
• The new unit tests were run as standalone pytest-style functions (no weights / GPU required).

Notes

• The FP8 shared expert is dequantized to bf16 at load (W8A16 precision). It is small relative to the 256 routed experts, but a native fp8 shared-expert module would preserve the weight-bandwidth benefit — that is a separate enhancement, not needed to fix the crash.
• This covers the modelopt MIXED_PRECISION format. A compressed-tensors mixed-precision checkpoint (128x128 block-FP8 shared expert + NVFP4 experts, e.g. some Qwen3.6-35B-A3B releases) uses config_groups rather than quantized_layers and is not handled here.